package io.openaev.rest.vulnerability.service;

import static io.openaev.helper.StreamHelper.fromIterable;
import static io.openaev.utils.pagination.PaginationUtils.buildPaginationJPA;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.openaev.config.cache.LicenseCacheManager;
import io.openaev.database.model.*;
import io.openaev.database.repository.CweRepository;
import io.openaev.database.repository.VulnerabilityRepository;
import io.openaev.ee.Ee;
import io.openaev.rest.collector.service.CollectorService;
import io.openaev.rest.exception.ElementNotFoundException;
import io.openaev.rest.vulnerability.form.*;
import io.openaev.utils.pagination.SearchPaginationInput;
import jakarta.annotation.Resource;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotNull;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Slf4j
@Service
@RequiredArgsConstructor
public class VulnerabilityService {

  private static final String VULNERABILITY_NOT_FOUND_MSG = "Vulnerability not found with id: ";

  private final CollectorService collectorService;
  private final Ee eeService;

  private final VulnerabilityRepository vulnerabilityRepository;
  private final CweRepository cweRepository;
  private final LicenseCacheManager licenseCacheManager;

  @Resource protected ObjectMapper mapper;

  public Vulnerability createVulnerability(final @Valid VulnerabilityCreateInput input) {
    final Vulnerability vulnerability = new Vulnerability();
    if (eeService.isEnterpriseLicenseInactive(licenseCacheManager.getEnterpriseEditionInfo())) {
      input.setRemediation(null);
    }
    vulnerability.setUpdateAttributes(input);
    updateCweAssociations(vulnerability, input.getCwes());
    return vulnerabilityRepository.save(vulnerability);
  }

  private List<Vulnerability> batchUpsertVulnerabilities(
      List<VulnerabilityCreateInput> vulnerabilityInput) {
    // Extract external IDs
    Set<String> externalIds =
        vulnerabilityInput.stream()
            .map(VulnerabilityCreateInput::getExternalId)
            .collect(Collectors.toSet());

    // Batch fetch existing vulnerabilities
    Map<String, Vulnerability> existingVulnerabilitiesMap =
        getVulnerabilitiesByExternalIds(externalIds).stream()
            .collect(Collectors.toMap(Vulnerability::getExternalId, Function.identity()));

    // Process with pre-fetched data
    List<Vulnerability> vulnerabilities =
        vulnerabilityInput.stream()
            .map(
                cveInput -> {
                  Vulnerability vulnerability =
                      existingVulnerabilitiesMap.getOrDefault(
                          cveInput.getExternalId(), new Vulnerability());
                  vulnerability.setUpdateAttributes(cveInput);
                  updateCweAssociations(vulnerability, cveInput.getCwes());
                  return vulnerability;
                })
            .toList();

    return fromIterable(vulnerabilityRepository.saveAll(vulnerabilities));
  }

  private void updateCollectorStateFromVulnerabilityBulkInsertInput(
      Collector collector, @NotNull VulnerabilityBulkInsertInput inputs) {
    ObjectNode collectorNewState = mapper.createObjectNode();
    collectorNewState.put(
        "last_modified_date_fetched", inputs.getLastModifiedDateFetched().toString());
    collectorNewState.put("last_index", inputs.getLastIndex().toString());
    collectorNewState.put("initial_dataset_completed", inputs.getInitialDatasetCompleted());
    this.collectorService.updateCollectorState(collector, collectorNewState);
  }

  @Transactional(rollbackFor = Exception.class)
  public void bulkUpsertVulnerabilities(@NotNull VulnerabilityBulkInsertInput inputs) {
    Collector collector = this.collectorService.collector(inputs.getSourceIdentifier());

    List<Vulnerability> vulnerabilities =
        this.batchUpsertVulnerabilities(inputs.getVulnerabilities());
    this.updateCollectorStateFromVulnerabilityBulkInsertInput(collector, inputs);

    log.info(
        "Bulk upsert {} vulnerabilities with last modified date fetched: {}",
        vulnerabilities.size(),
        inputs.getLastModifiedDateFetched());
  }

  public Page<Vulnerability> searchVulnerabilities(final @Valid SearchPaginationInput input) {
    return buildPaginationJPA(
        (Specification<Vulnerability> spec, Pageable pageable) ->
            vulnerabilityRepository.findAll(spec, pageable),
        input,
        Vulnerability.class);
  }

  public Vulnerability updateVulnerability(
      final String vulnerabilityId, final @Valid VulnerabilityUpdateInput input) {
    final Vulnerability existingVulnerability = findById(vulnerabilityId);
    if (eeService.isEnterpriseLicenseInactive(licenseCacheManager.getEnterpriseEditionInfo())) {
      input.setRemediation(null);
      BeanUtils.copyProperties(input, existingVulnerability, "remediation");
    } else {
      existingVulnerability.setUpdateAttributes(input);
    }
    updateCweAssociations(existingVulnerability, input.getCwes());
    return vulnerabilityRepository.save(existingVulnerability);
  }

  public Vulnerability findById(final String vulnerabilityId) {
    return vulnerabilityRepository
        .findById(vulnerabilityId)
        .orElseThrow(
            () -> new ElementNotFoundException(VULNERABILITY_NOT_FOUND_MSG + vulnerabilityId));
  }

  public Set<Vulnerability> findAllByIdsOrThrowIfMissing(final Set<String> vulnIds) {
    Set<Vulnerability> vulns = this.vulnerabilityRepository.getAllByIdInIgnoreCase(vulnIds);
    throwIfMissing(vulnIds, vulns, Vulnerability::getId);
    return vulns;
  }

  public Vulnerability findByExternalId(String externalId) {
    return vulnerabilityRepository
        .findByExternalId(externalId)
        .orElseThrow(() -> new ElementNotFoundException(VULNERABILITY_NOT_FOUND_MSG + externalId));
  }

  public Set<Vulnerability> findAllByExternalIdsOrThrowIfMissing(final Set<String> vulnIds) {
    Set<Vulnerability> vulns = getVulnerabilitiesByExternalIds(vulnIds);
    throwIfMissing(vulnIds, vulns, Vulnerability::getExternalId);
    return vulns;
  }

  private void throwIfMissing(
      Set<String> requiredIds,
      Set<Vulnerability> fetchedVulnerabilities,
      Function<? super Vulnerability, String> getId) {

    List<String> fetchedIdLower =
        fetchedVulnerabilities.stream().map(vuln -> getId.apply(vuln).toLowerCase()).toList();

    List<String> missingIds =
        requiredIds.stream().filter(id -> !fetchedIdLower.contains(id.toLowerCase())).toList();

    if (!missingIds.isEmpty()) {
      throw new ElementNotFoundException(
          String.format("Missing vulnerabilities: %s", String.join(", ", missingIds)));
    }
  }

  public void deleteById(final String vulnerabilityId) {
    vulnerabilityRepository.deleteById(vulnerabilityId);
  }

  private void updateCweAssociations(Vulnerability vulnerability, List<CweInput> cweInputs) {
    if (cweInputs == null || cweInputs.isEmpty()) {
      vulnerability.setCwes(Collections.emptyList());
      return;
    }

    List<Cwe> cweEntities =
        cweInputs.stream()
            .map(
                input ->
                    cweRepository
                        .findByExternalId(input.getExternalId())
                        .orElseGet(
                            () -> {
                              Cwe newCwe = new Cwe();
                              newCwe.setExternalId(input.getExternalId());
                              newCwe.setSource(input.getSource());
                              return cweRepository.save(newCwe);
                            }))
            .collect(Collectors.toList());

    vulnerability.setCwes(cweEntities);
  }

  /**
   * Resolves external Vulnerability Refs from a set of vulnerability {@link Vulnerability}
   * entities.
   *
   * @param externalIds set vulnerability Refs
   * @return set of resolved vulnerability entities
   */
  public Set<Vulnerability> getVulnerabilitiesByExternalIds(Set<String> externalIds) {
    if (externalIds.isEmpty()) {
      return Collections.emptySet();
    }
    return this.vulnerabilityRepository.getAllByExternalIdInIgnoreCase(externalIds);
  }
}
